import torch
import torch.nn.functional as F
import itertools
from utils.image_pool import ImagePool
from .base_model import BaseModel
from .cyclegan import networks
from .cyclegan.testapi import TestApi,get_frequency

import numpy as np

# from .lightcnn import network as lightcnn
# from models import create_model
# from copy import deepcopy
# import mypymath
# import time
# import os


# def recoverFslope(slope):
#     B,_,H,W = slope.shape
#     # x = slope[:,0:1,:,:]
#     # x1 = recove_h1_slope(slope[:,1:2,:,:])
#     # x2 = recove_v2_slope(slope[:,2:3:,:,:])
#     # x3 = recove_ldiag3_slope(slope[:,3:4,:,:])
#     # x4 = recove_rdiag4_slope(slope[:,4:5,:,:])
#     x = deepcopy(slope[:,0:1,:,:])
#     x1 = deepcopy(slope[:,1:2,:,:])
#     x2 = deepcopy(slope[:,2:3:,:,:])
#     recove_h1_slope(x1)
#     recove_v2_slope(x2)
    
#     if slope.shape[1] == 5:
#         x3 = deepcopy(slope[:,3:4,:,:])
#         x4 = deepcopy(slope[:,4:5,:,:])
#         recove_ldiag3_slope(x3)
#         recove_rdiag4_slope(x4)
#         return torch.mean(torch.cat((x,x1,x2,x3,x4,),dim=1),dim=1).reshape(B,1,H,W)
#     return torch.mean(torch.cat((x,x1,x2,),dim=1),dim=1).reshape(B,1,H,W)
#     pass
        
# def recove_h1_slope(slope_img):
#     # print("=="*20)
#     # print("mydata:")
#     # print(data[0,0,:5,:5])
#     # ans = torch.zeros(slope_img.shape).to(slope_img.device)
#     # ans[:,:,:,0] = slope_img[:,:,:,0]
#     for i in range(1,slope_img.shape[-1]):
#         # ans[:,:,:,i] = ans[:,:,:,i-1]+slope_img[:,:,:,i]
#         slope_img[:,:,:,i]+=slope_img[:,:,:,i-1]
#     return slope_img
    
# def recove_v2_slope(slope_img):
#     # print("=="*20)
#     # print("mydata:")
#     # print(data[0,0,:5,:5])
#     # ans = torch.zeros(slope_img.shape).to(slope_img.device)
#     # ans[:,:,0,:] = slope_img[:,:,0,:]
#     for i in range(1,slope_img.shape[-1]):
#         # ans[:,:,i,:] = ans[:,:,i-1,:]+slope_img[:,:,i,:]
#         slope_img[:,:,i,:]+=slope_img[:,:,i-1,:]
#     return slope_img

# def recove_ldiag3_slope(slope_img):
#     # print("=="*20)
#     # print("mydata:")
#     # print(data[0,0,:5,:5])
#     # ans = torch.zeros(slope_img.shape).to(slope_img.device)
#     # ans[:,:,:,0] = slope_img[:,:,:,0] 
#     # ans[:,:,0,:] = slope_img[:,:,0,:]
#     # for j in range(1,slope_img.shape[-1]):
#     #     for i in range(1,slope_img.shape[-2]):
#     #         ans[:,:,i,j] = ans[:,:,i-1,j-1]+slope_img[:,:,i,j]
#     # return ans
#     for i in range(1,slope_img.shape[-1]):
#         for j in range(1,slope_img.shape[-2]):
#         # ans[:,:,i,:] = ans[:,:,i-1,:]+slope_img[:,:,i,:]
#             slope_img[:,:,i,j] += slope_img[:,:,i-1,j-1]
#     return slope_img

# def recove_rdiag4_slope(slope_img):
#     # print("=="*20)
#     # print("mydata:")
#     # print(data[0,0,:5,:5])
#     # ans = torch.zeros(slope_img.shape).to(slope_img.device)
#     # ans[:,:,:,slope_img.shape[-1]-1] = slope_img[:,:,:,slope_img.shape[-1]-1]
#     # ans[:,:,0,:] = slope_img[:,:,0,:]
#     for j in range(slope_img.shape[-1]-2,-1,-1):
#         for i in range(1,slope_img.shape[-2]):
#             slope_img[:,:,i,j] += slope_img[:,:,i-1,j+1]
#     return slope_img
    
class CycleGANS012DFFTModel(BaseModel):
    """
    This class implements the CycleGAN model, for learning image-to-image translation without paired data.

    The model training requires '--dataset_mode unaligned' dataset.
    By default, it uses a '--netG resnet_9blocks' ResNet generator,
    a '--netD basic' discriminator (PatchGAN introduced by pix2pix),
    and a least-square GANs objective ('--gan_mode lsgan').

    CycleGAN paper: https://arxiv.org/pdf/1703.10593.pdf
    """
    @staticmethod
    def modify_commandline_options(parser, is_train=True):
        """Add new dataset-specific options, and rewrite default values for existing options.

        Parameters:
            parser          -- original option parser
            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
        
        Returns:
            the modified parser.

        For CycleGAN, in addition to GAN losses, we introduce lambda_A, lambda_B, and lambda_identity for the following losses.
        A (source domain), B (target domain).
        Generators: G_A: A -> B; G_B: B -> A.
        Discriminators: D_A: G_A(A) vs. B; D_B: G_B(B) vs. A.
        Forward cycle loss:  lambda_A * ||G_B(G_A(A)) - A|| (Eqn. (2) in the paper)
        Backward cycle loss: lambda_B * ||G_A(G_B(B)) - B|| (Eqn. (2) in the paper)
        Identity loss (optional): lambda_identity * (||G_A(B) - B|| * lambda_B + ||G_B(A) - A|| * lambda_A) (Sec 5.2 "Photo generation from paintings" in the paper)
        Dropout is not used in the original CycleGAN paper.
        """
        parser.set_defaults(no_dropout=True)  # default CycleGAN did not use dropout
        parser.set_defaults(frequency=True)
        # parser.add_argument('--num_classes', type=int, default=3, help='the architecture situation of model')
        # parser.add_argument('--pretrained_model', type=str, default="lightcnn", help='the architecture of pretrained model')
        # parser.add_argument('--pretrained_path', type=str, default="lightcnn", help='the path of pretrained model')
        # parser.add_argument('--pretrained_ndf', type=int, default=128, help='the number of features of output ')
        # parser.add_argument('--pretrained_output_nc', type=int, default=256, help='the number of features of output ')
        # parser.add_argument('--bins', type=int, default=16, help='the number of features of output ')
        parser.add_argument('--cal_mode', type=str, default='linear', help='the mode of caculating')
        parser.add_argument('--D_generation', type=str, default='mergeS12', help='mergeS12 or source')
        parser.add_argument('--slope', type=int, default=0, help='slope parameter')
        parser.add_argument('--use_inception',action='store_true', help='The switch of inception')
        # parser.add_argument('--frequency', type=bool, default=True, help='the mode of caculating')
        # if is_train:
        parser.add_argument('--lambda_A', type=float, default=1.0, help='weight for cycle loss (A -> B -> A)')
        parser.add_argument('--lambda_B', type=float, default=1.0, help='weight for cycle loss (B -> A -> B)')
        parser.add_argument('--lambda_D', type=float, default=1.0, help='weight for Adversarial network loss')
        parser.add_argument('--lambda_fre', type=float, default=0.02, help='weight for Frequency loss')
        parser.add_argument('--lambda_idt', type=float, default=0.5, help='use identity mapping. Setting lambda_identity other than 0 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set lambda_identity = 0.1')
            

        return parser

    def __init__(self, opt):
        """Initialize the CycleGAN class.

        Parameters:
            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
        """
        
        BaseModel.__init__(self, opt)
        # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
        # self.loss_names = ['D_A', 'G_A', 'cycle_A', 'idt_A', 'D_B', 'G_B', 'cycle_B', 'idt_B']
        # self.loss_chage_line = "\n"
        self.opt.frequency = not self.opt.lambda_fre == 0
        self.TestApi = TestApi(opt)
        self.visuale_param()

        # define networks (both Generators and discriminators)
        # The naming is different from those used in the paper.
        # Code (vs. paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X)
        self.netG_S01 = networks.define_G(opt.input_nc + opt.slope, opt.output_nc, opt.ngf, opt.netG, opt.norm,
                                        not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,inception=opt.use_inception)
        
        self.netG_S10 = networks.define_G(opt.output_nc + opt.slope, opt.input_nc, opt.ngf, opt.netG, opt.norm,
                                        not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,inception=opt.use_inception)
        
        
        self.netG_S02 = networks.define_G(opt.input_nc + opt.slope, opt.output_nc, opt.ngf, opt.netG, opt.norm,
                                        not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,inception=opt.use_inception)
        
        self.netG_S20 = networks.define_G(opt.output_nc + opt.slope, opt.input_nc, opt.ngf, opt.netG, opt.norm,
                                        not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,inception=opt.use_inception)
        
        if self.opt.D_generation == 'source':
            self.netG_S0D = networks.define_G(opt.input_nc + opt.slope, opt.output_nc, opt.ngf, opt.netG, opt.norm,
                                            not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,inception=opt.use_inception)
            
            self.netG_SD0 = networks.define_G(opt.output_nc + opt.slope, opt.input_nc, opt.ngf, opt.netG, opt.norm,
                                            not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,inception=opt.use_inception)


        if self.isTrain or self.opt.use_inception:  # define discriminators
            self.netD_S1 = networks.define_D(opt.output_nc, opt.ndf, opt.netD,
                                            opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)

            self.netD_S0 = networks.define_D(opt.input_nc, opt.ndf, opt.netD,
                                            opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
            
            self.netD_S2 = networks.define_D(opt.output_nc, opt.ndf, opt.netD,
                                            opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
                
            self.netD_SD = networks.define_D(opt.input_nc, opt.ndf, opt.netD,
                                            opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
            # self.netD_light.define_lightcnn(opt)

            # self.netD_S02 = networks.define_D(opt.input_nc, opt.ndf, opt.netD,
            #                                 opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
        self.criterionPredD = torch.nn.L1Loss()
        
        self.opt.lambda_idt = self.opt.lambda_idt if self.opt.lambda_idt > 0 else 0
            # define loss functions
        self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)  # define GAN loss.
        self.criterionCycle = torch.nn.L1Loss()
        self.criterionIdt = torch.nn.L1Loss()
        
        if self.isTrain:
            if self.opt.continue_train:
                self.epoch = self.opt.epoch_count
            # self.nepoch = self.opt.n_epochs + self.opt.n_epochs_decay
            # self.fac = self.epoch/(self.nepoch+1)
            if opt.lambda_idt > 0.0:  # only works when input and output images have the same number of channels
                assert(opt.input_nc == opt.output_nc)
            self.fake_S0_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
            self.fake_S1_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
            self.fake_S2_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
            self.fake_SD_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
            # self.fake_S2_S0_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
            
        # if self.isTrain:
            # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
            if self.opt.D_generation == 'source':
                self.optimizer_G = torch.optim.Adam(itertools.chain(
                                                                self.netG_S01.parameters(),
                                                                self.netG_S10.parameters(),
                                                                self.netG_S02.parameters(),
                                                                self.netG_S20.parameters(),
                                                                self.netG_S0D.parameters(),
                                                                self.netG_SD0.parameters(),
                                                                ),
                                                    lr=opt.lr, betas=(opt.beta1, 0.999))
            elif self.opt.D_generation == 'mergeS12':
                self.optimizer_G = torch.optim.Adam(itertools.chain(
                                                                self.netG_S01.parameters(),
                                                                self.netG_S10.parameters(),
                                                                self.netG_S02.parameters(),
                                                                self.netG_S20.parameters(),
                                                                ),
                                                    lr=opt.lr, betas=(opt.beta1, 0.999))

            self.optimizer_D = torch.optim.Adam(itertools.chain(
                                                                self.netD_S0.parameters(),
                                                                self.netD_S1.parameters(),
                                                                self.netD_S2.parameters(),
                                                                self.netD_SD.parameters(),
                                                                ),
                                                lr=opt.lr, betas=(opt.beta1, 0.999))

            self.optimizers.append(self.optimizer_G)
            self.optimizers.append(self.optimizer_D)
    
    def visuale_param(self):
        
        # if not "S1_SSIM" in self.loss_names:
        self.TestApi.init(self.isTrain)
        if self.isTrain:
            self.loss_names = ["mean_G_S10","mean_G_S20","mean_G_SD0",
                            "mean_G_S01","mean_G_S02","mean_G_S0D",
                            "mean_G",
                            ]
            self.loss_names += ["mean_D_S0","mean_D_S1","mean_D_S2","mean_D_SD"]
            
        else:
            # self.loss_names = ["S1_SSIM","S01_SSIM","S2_SSIM","S02_SSIM"]
            # self.loss_names += ["PSNR_S1","S01_PSNR","PSNR_S2","S02_PSNR",]
            self.loss_names = ["mean_PSNR_S1","mean_S01_PSNR","mean_PSNR_S2","mean_S02_PSNR"]
            self.loss_names += ["mean_SSIM_S1","mean_S01_SSIM","mean_SSIM_S2","mean_S02_SSIM"]
            self.loss_names += ["mean_FID_S1","mean_FID_S01","mean_DFID_S1"]
            self.loss_names += ["mean_FID_S2","mean_FID_S02","mean_DFID_S2"]
            
            if self.opt.D_generation == "source":
                # self.loss_names += ["PSNR_SD","S0D_PSNR","SD_SSIM","S0D_SSIM",]
                self.loss_names += ["mean_SSIM_SD","mean_S0D_SSIM"]
                self.loss_names += ["mean_FID_SD","mean_FID_S0D","mean_DFID_SD"]
                # if not self.isTrain:
                self.loss_names += ["mean_PSNR_SD","mean_S0D_PSNR"]
                
            elif self.opt.D_generation == "mergeS12":
                # self.loss_names += ["PSNR_S12D","S0_12D_PSNR","SSIM_S12D","S0_12D_SSIM"]
                self.loss_names += ["mean_SSIM_S12D","mean_S0D_SSIM"]
                # if not self.isTrain:
                self.loss_names += ["mean_PSNR_S12D","mean_S0D_PSNR",]
                self.loss_names += ["mean_FID_S12D","mean_FID_S0D","mean_DFID_SD"]
                
                
            # self.loss_names += ["TOP10p_mean_S012D_SSIM,TOP50p_mean_S012D_SSIM,TOP100p_mean_S012D_SSIM"]
            # self.count_mean_S012D_SSIM=NewAverageMeter()
        if self.opt.frequency:
            self.loss_names += ["mean_F_S01","mean_F_S02","mean_F_S0D",
                                "mean_F_S1","mean_F_S2","mean_F_SD",
                                "mean_F_S12D"]
        # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
        visual_names_S1A = []
        visual_names_S1B = []
        visual_names_S1A += ['real_S0', 'fake_S01', 'rec_S10']
        # visual_names_S1A += ['real_S0_h', 'fake_S01_h', 'rec_S10_h']
        visual_names_S1A += ['fake_S02', 'rec_S20']
        # visual_names_S1A += [ 'fake_S02_h', 'rec_S20_h']
        visual_names_S1A += [ 'fake_S0D', 'rec_SD0']
        # visual_names_S1A += [ 'fake_S0D_h', 'rec_SD0_h']
        visual_names_S1A += ['real_S1', 'real_S2', 'real_SD', 'fake_S12D']
        # visual_names_S1B += [, 'real_S12D',]
        # visual_names_S1A += ['real_S2_h', 'real_SD_h', 'real_S1_h']
        
        if self.isTrain:
            visual_names_S1B += ['fake_S10', 'rec_S01']
            # visual_names_S1B += ['fake_S10_h', 'rec_S01_h']
            visual_names_S1B += ['fake_S20', 'rec_S02']
            # visual_names_S1B += ['fake_S20_h', 'rec_S02_h']
            visual_names_S1B += ['fake_SD0', 'rec_S0D']
            # visual_names_S1B += ['fake_SD0_h', 'rec_S0D_h']
        
        self.visual_names = visual_names_S1A + visual_names_S1B  # combine visualizations for A and B
        # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>.
        self.set_model_name()
    
    def set_model_name(self):
        if self.isTrain or self.opt.use_inception:
            # self.model_names = ['G_A', 'G_B', 'D_A', 'D_B']
            self.model_names = ['G_S01', 'G_S10', 
                                'G_S02', 'G_S20', 
                                'D_S0', 'D_S1', 'D_S2','D_SD',]
            if self.opt.D_generation == "source":
                self.model_names += ['G_S0D', 'G_SD0',]
        else:  # during test time, only load Gs
            self.model_names = ['G_S01', 'G_S10','G_S02', 'G_S20', ]
            if self.opt.D_generation == "source":
                self.model_names += ['G_S0D', 'G_SD0',]
        return
    # def cal_slope(self,data,directions=2,mode='constant',):
    #     if directions == 0:
    #         return data.to(self.device)
    #     x = F.pad(data,(1,1,1,1,0,0),mode,value=0)
    #     x1 = x[:,:,1:-1,1:-1]-x[:,:,1:-1,0:-2]
    #     x2 = x[:,:,1:-1,1:-1]-x[:,:,0:-2,1:-1]
    #     if directions == 2:
    #         return torch.cat((x[:,:,1:-1,1:-1],x1,x2),dim=1).to(self.device)
    #     if directions == 4:
    #         x3 = x[:,:,1:-1,1:-1]-x[:,:,0:-2,0:-2]
    #         x4 = x[:,:,1:-1,1:-1] - x[:,:,0:-2,2:]
    #         # test = torch.cat((x[:,:,1:-1,1:-1],x1,x2,x3,x4),dim=1).to(self.device)
    #         # recoverFslope(test)
    #         return torch.cat((x[:,:,1:-1,1:-1],x1,x2,x3,x4),dim=1).to(self.device)
            
        
    def set_input(self, input):
        """Unpack input data from the dataloader and perform necessary pre-processing steps.

        Parameters:
            input (dict): include the data itself and its metadata information.

        The option 'direction' can be used to swap domain A and domain B.
        """
        # AtoB = self.opt.direction == 'AtoB'
        # self.real_A = input['A' if AtoB else 'B'].to(self.device)
        # self.real_B = input['B' if AtoB else 'A']self.device)
        # self.image_paths = input['A_paths' if AtoB else 'B_paths']
        # if "isTrain" in input.keys():
        visual_param_changed = False
        if "isTrain" in input.keys():
            if not self.isTrain == input["isTrain"]:
                self.isTrain = input["isTrain"]
                visual_param_changed=True
                
        if 'frequency' in input.keys() and not self.opt.frequency == input["frequency"][0].item():
            self.opt.frequency = input["frequency"][0].item()
            visual_param_changed = True
        
        if visual_param_changed:
            self.visuale_param()
        
        self.image_paths = input['S0_paths']
        # if self.isTrain:
        self.real_S0 = input['S0'].to(self.device)
        self.real_S1 = input['S1'].to(self.device)
        self.real_S2 = input['S2'].to(self.device)
        self.real_SD = input['SD'].to(self.device)
        # self.real_S0_ori = input['S0']
        # self.real_S1_ori = input['S1']
        # self.real_S2_ori = input['S2']
        # self.real_SD_ori = input['SD']
        
        # self.real_S0 = self.cal_slope(input['S0'],self.opt.slope)
        # self.real_S1 = self.cal_slope(input['S1'],self.opt.slope)
        # self.real_S2 = self.cal_slope(input['S2'],self.opt.slope)
        # self.real_SD = self.cal_slope(input['SD'],self.opt.slope)
        
        
        # if self.opt.slope > 0:
        #     self.real_S0_slope = 
        #     self.real_S1_slope = self.cal_slope(input['S1'],self.opt.slope)
        #     self.real_S2_slope = self.cal_slope(input['S2'],self.opt.slope)
        #     self.real_SD_slope = self.cal_slope(input['SD'],self.opt.slope)
        # self.real_S0 = normal(self.real_S0,1,1e-5)
        # self.real_S1 = normal(self.real_S1,1,1e-5)
        # self.real_S2 = normal(self.real_S2,1,1e-5)
        # self.real_SD = normal(self.real_SD,1,1e-5)
            # self.real_S2 = input['S2'].to(self.device)
            
        # self.real_S12D = torch.sqrt(self.real_S1**2 + self.real_S2**2)/(self.real_S0+1e-5)
        # self.real_S12D = normal(self.real_S12D,1,1e-5)
                
        if self.opt.frequency:
            self.real_S0_F = input['S0_F'][:,0,:,:]
            self.real_S1_F = input['S1_F'][:,0,:,:]
            self.real_S2_F = input['S2_F'][:,0,:,:]
            self.real_SD_F = input['SD_F'][:,0,:,:]
        

    def forward(self):
        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
        # self.fake_B = self.netG_A(self.real_A)  # G_A(A)
        # self.rec_A = self.netG_B(self.fake_B)   # G_B(G_A(A))
        # self.fake_A = self.netG_B(self.real_B)  # G_B(B)
        # self.rec_B = self.netG_A(self.fake_A)   # G_A(G_B(B))

        self.fake_S01 = self.netG_S01(self.real_S0, self.opt.slope)
        self.rec_S10 = self.netG_S10(self.fake_S01["generation"], self.opt.slope)
        
        self.fake_S02 = self.netG_S02(self.real_S0, self.opt.slope)
        self.rec_S20 = self.netG_S20(self.fake_S02["generation"], self.opt.slope)
        
        # if self.opt.slope > 0 and self.opt.frequency:
        #     self.fake_S01_recover =  recoverFslope(self.fake_S01, self.opt.slope)
        #     self.rec_S10_recover =  recoverFslope(self.rec_S10, self.opt.slope)
        #     self.fake_S02_recover =  recoverFslope(self.fake_S02, self.opt.slope)
        #     self.rec_S20_recover =  recoverFslope(self.rec_S20, self.opt.slope)
        
        if self.opt.D_generation == "source":
            self.fake_S0D = self.netG_S0D(self.real_S0, self.opt.slope)
            self.rec_SD0 = self.netG_SD0(self.fake_S0D["generation"], self.opt.slope)
            # if self.opt.slope > 0 and self.opt.frequency:
                # self.fake_S0D_recover =  recoverFslope(self.fake_S0D)
                # self.rec_SD0_recover =  recoverFslope(self.rec_SD0)
        
        if self.isTrain:
            
            self.fake_S10 = self.netG_S10(self.real_S1, self.opt.slope)
            self.rec_S01 = self.netG_S01(self.fake_S10["generation"], self.opt.slope)
            
            self.fake_S20 = self.netG_S20(self.real_S2, self.opt.slope)
            self.rec_S02 = self.netG_S02(self.fake_S20["generation"], self.opt.slope)
            
            # if self.opt.slope > 0 and self.opt.frequency:
                # self.fake_S10_recover =  recoverFslope(self.fake_S10)
                # self.rec_S01_recover =  recoverFslope(self.rec_S01)
                # self.fake_S20_recover =  recoverFslope(self.fake_S20)
                # self.rec_S02_recover =  recoverFslope(self.rec_S02)
            
            if self.opt.D_generation == "source":
                self.fake_SD0 = self.netG_SD0(self.real_SD, self.opt.slope)
                self.rec_S0D = self.netG_S0D(self.fake_SD0["generation"], self.opt.slope)
                # if self.opt.slope > 0 and self.opt.frequency:
                    # self.fake_SD0_recover =  recoverFslope(self.fake_SD0)
                    # self.rec_S0D_recover =  recoverFslope(self.rec_S0D)
            
        # self.fake_S01 = normal(self.fake_S01,1,1e-5)
        # self.fake_S02 = normal(self.fake_S02,1,1e-5)
        # x1 = self.fake_S01/self.real_S0
        # x2 = self.fake_S02/self.real_S0
        # x1 = normal(x1,1,-1)
        # x2 = normal(x2,1,-1)
            elif self.opt.D_generation == "mergeS12":
                self.fake_S12D = {}
                self.fake_S12D["generation"] = torch.sqrt(self.fake_S01["generation"]**2 + self.fake_S02["generation"]**2)
            # self.fake_S12D = normal(self.fake_S12D,1,1e-5)
        else:
            if self.opt.use_inception:
                real_S01_Fea = self.netD_S1(self.real_S0)
                real_S02_Fea = self.netD_S2(self.real_S0)
                real_S0D_Fea = self.netD_SD(self.real_S0)
                
                real_S1_Fea = self.netD_S1(self.real_S1)
                real_S2_Fea = self.netD_S2(self.real_S2)
                real_SD_Fea = self.netD_SD(self.real_SD)
                fake_S1_Fea = self.netD_S1(self.fake_S01["generation"])
                fake_S2_Fea = self.netD_S2(self.fake_S02["generation"])
                if self.opt.D_generation == "source":
                    fake_SD_Fea = self.netD_SD(self.fake_S0D["generation"])
                elif self.opt.D_generation == "mergeS12":
                    fake_SD_Fea = self.netD_SD(self.fake_S12D["generation"])
                
                self.loss_FID_S1 = self.criterionGAN.loss(real_S1_Fea,fake_S1_Fea)
                self.loss_FID_S2 = self.criterionGAN.loss(real_S2_Fea,fake_S2_Fea)
                self.loss_FID_SD = self.criterionGAN.loss(real_SD_Fea,fake_SD_Fea)
                self.loss_FID_S01 = self.criterionGAN.loss(real_S01_Fea,real_S1_Fea)
                self.loss_FID_S02 = self.criterionGAN.loss(real_S02_Fea,real_S2_Fea)
                self.loss_FID_S0D = self.criterionGAN.loss(real_S0D_Fea,real_SD_Fea)
            pass
        
    def backward_D_basic(self, netD, real, fake):
        """Calculate GAN loss for the discriminator

        Parameters:
            netD (network)      -- the discriminator D
            real (tensor array) -- real images
            fake (tensor array) -- images generated by a generator

        Return the discriminator loss.
        We also call loss_D.backward() to calculate the gradients.
        """
        # Real
        pred_real = netD(real)
        loss_D_real = self.criterionGAN(pred_real, True)
        # Fake
        pred_fake = netD(fake.detach())
        loss_D_fake = self.criterionGAN(pred_fake, False)
        # Combined loss and calculate gradients
        loss_D = (loss_D_real + loss_D_fake) * 0.5
        if self.isTrain:
            loss_D.backward()
        return loss_D

    def backward_D_S0(self):
        """Calculate GAN loss for discriminator D_S0"""
        if self.opt.D_generation == "source":
            fake_S0 = torch.cat([self.fake_S10["generation"],self.fake_S20["generation"],self.fake_SD0["generation"]])
        elif self.opt.D_generation == "mergeS12":
            fake_S0 = torch.cat([self.fake_S10["generation"],self.fake_S20["generation"]])
        fake_S0 = self.fake_S0_pool.query(fake_S0)
        self.loss_D_S0 = self.backward_D_basic(self.netD_S0, self.real_S0, fake_S0)
        
    def backward_D_S1(self):
        """Calculate GAN loss for discriminator D_S1"""
        fake_S01 = self.fake_S1_pool.query(self.fake_S01["generation"])
        self.loss_D_S1 = self.backward_D_basic(self.netD_S1, self.real_S1, fake_S01)
        
    def backward_D_S2(self):
        """Calculate GAN loss for discriminator D_S1"""
        fake_S02 = self.fake_S2_pool.query(self.fake_S02["generation"])
        self.loss_D_S2 = self.backward_D_basic(self.netD_S2, self.real_S2, fake_S02)
        
    def backward_D_SD(self):
        """Calculate GAN loss for discriminator D_S1"""
        if self.opt.D_generation == "source":
            fake_S0D = self.fake_S1_pool.query(self.fake_S0D["generation"])
        elif self.opt.D_generation == "mergeS12":
            fake_S0D = self.fake_S1_pool.query(self.fake_S12D["generation"])
            
        # fake_S0D = self.fake_S1_pool.query(torch.cat([self.fake_S0D,self.fake_S12D]))
        self.loss_D_SD = self.backward_D_basic(self.netD_SD, self.real_SD, fake_S0D)
        
    
    def set_Frequency_S01(self,loss1,loss2=0):
        self.loss_F_S1 = (loss1 + loss2).to(self.device)
    
    def set_Frequency_S10(self,loss1,loss2=0):
        self.loss_F_S10 = (loss1 + loss2).to(self.device)
        
    def set_Frequency_S02(self,loss1,loss2=0):
        self.loss_F_S2 = (loss1 + loss2).to(self.device)
        
    def set_Frequency_S20(self,loss1,loss2=0):
        self.loss_F_S20 = (loss1 + loss2).to(self.device)
    
    def set_Frequency_S0D(self,loss1,loss2=0):
        self.loss_F_SD = (loss1 + loss2).to(self.device)
    
    def set_Frequency_SD0(self,loss1,loss2=0):
        self.loss_F_SD0 = (loss1 + loss2).to(self.device)
    
    def backward_Frequency(self,X):
        
        # starttime1 = time.time()
        # fake_S0X_F = get_frequency(getattr(self,'fake_S0'+X)["generation"].contiguous().cpu().detach().numpy()[:,0,:,:])
        # rec_SX0_F = get_frequency(getattr(self,'rec_S'+X+'0')["generation"].contiguous().cpu().detach().numpy()[:,0,:,:])
        # getattr(self,'set_Frequency_S0'+X)(self.criterionPredD(fake_S0X_F,getattr(self,'real_S'+X+'_F')))
        # getattr(self,'set_Frequency_S'+X+'0')(self.criterionPredD(rec_SX0_F,getattr(self,'real_S'+X+'_F')))
        if self.isTrain:
            rec_S0X_F = get_frequency(getattr(self,'rec_S0'+X)["generation"].contiguous().cpu().detach().numpy()[:,0,:,:])
            fake_SX0_F = get_frequency(getattr(self,'fake_S'+X+'0')["generation"].contiguous().cpu().detach().numpy()[:,0,:,:])
            getattr(self,'set_Frequency_S0'+X)(self.criterionPredD(fake_S0X_F,getattr(self,'real_S'+X+'_F')),self.criterionPredD(rec_S0X_F,getattr(self,'real_S'+X+'_F')))
            getattr(self,'set_Frequency_S'+X+'0')(self.criterionPredD(fake_SX0_F,getattr(self,'real_S0_F')),self.criterionPredD(rec_SX0_F,getattr(self,'real_S0_F')))
        else:
            fake_S0X_F = get_frequency(getattr(self,'fake_S0'+X)["generation"].contiguous().cpu().detach().numpy()[:,0,:,:])
            rec_SX0_F = get_frequency(getattr(self,'rec_S'+X+'0')["generation"].contiguous().cpu().detach().numpy()[:,0,:,:])
            getattr(self,'set_Frequency_S0'+X)(self.criterionPredD(fake_S0X_F,getattr(self,'real_S'+X+'_F')))
            getattr(self,'set_Frequency_S'+X+'0')(self.criterionPredD(rec_SX0_F,getattr(self,'real_S'+X+'_F')))
        # getattr(self,'set_Frequency_rec_S0'+'X')(self.criterionPredD(rec_S0X_F,getattr(self,'real_S'+X+'_F')),self.criterionPredD(rec_S0X_F,getattr(self,'real_S'+X+'_F')))
        # getattr(self,'set_Frequency_rec_S'+X+'0')(self.criterionPredD(rec_SX0_F,getattr(self,'real_S0_F')),self.criterionPredD(rec_S0X_F,getattr(self,'real_S'+X+'_F')))
        # getattr(self,'loss_Frequency'+) = self.loss_Frequency_fake_S0 + self.loss_Frequency_fake_S1 + self.loss_Frequency_rec_S0 + self.loss_Frequency_rec_S1
        # getattr(self,count_Frequency).update(self.loss_Frequency.item())
        # getattr(self,loss_mean_Frequency) = self.count_Frequency.avg

    def cal_param_G_S01(self):
        self.loss_G_S01 = torch.tensor(float(0)).to(self.device)
        # if not self.opt.use_inception:
        self.idt_S01 = self.netG_S01(self.real_S1,self.opt.slope)
        # else:
        #     idt_S01,_ = self.netG_S01(self.real_S1,self.opt.slope)
        loss_idt_S01 = self.criterionIdt(self.idt_S01["generation"], self.real_S1) * self.opt.lambda_B * self.opt.lambda_idt
        loss_cycle_S01 = (self.criterionCycle(self.rec_S01, self.real_S1) )* self.opt.lambda_B
        
        self.loss_G_S01 += loss_idt_S01 +loss_cycle_S01
        if self.opt.frequency:
            self.loss_G_S01 += self.loss_F_S1 * self.opt.lambda_fre
        if self.isTrain:
            self.loss_G_S01 += self.criterionGAN(self.netD_S1(self.fake_S01["generation"]), True)*self.opt.lambda_D
            # self.loss_G_S01 += self.criterionCycle(self.fake_S01, self.real_S1)
        pass
    
    def cal_param_G_S10(self):
        self.loss_G_S10 = torch.tensor(float(0)).to(self.device)
        # if not self.opt.use_inception:
        idt_S10= self.netG_S10(self.real_S0,self.opt.slope)
        # else:
        #     idt_S10,_ = self.netG_S10(self.real_S0,self.opt.slope)
        loss_idt_S10 = self.criterionIdt(idt_S10["generation"], self.real_S0) * self.opt.lambda_A * self.opt.lambda_idt
        loss_cycle_S10 = self.criterionCycle(self.rec_S10, self.real_S0) * self.opt.lambda_A
        
        self.loss_G_S10 += loss_cycle_S10 + loss_idt_S10
        if self.opt.frequency:
            self.loss_G_S10 += self.loss_F_S10 * self.opt.lambda_fre
        if self.isTrain:
            self.loss_G_S10 += self.criterionGAN(self.netD_S0(self.fake_S10["generation"]), True)*self.opt.lambda_D
            # self.loss_G_S10 += self.criterionCycle(self.fake_S10, self.real_S0)
            # self.loss_G_S10.backward()
        pass
    
    def cal_param_G_S02(self):
        self.loss_G_S02 = torch.tensor(float(0)).to(self.device)
        # if not self.opt.use_inception:
        self.idt_S02 = self.netG_S02(self.real_S2,self.opt.slope)
        # else:
        #     idt_S02,_ = self.netG_S02(self.real_S2,self.opt.slope)
        loss_idt_S02 = self.criterionIdt(self.idt_S02["generation"], self.real_S2) * self.opt.lambda_B * self.opt.lambda_idt
        loss_cycle_S02 = self.criterionCycle(self.rec_S02, self.real_S2) * self.opt.lambda_B
        
        self.loss_G_S02 += loss_idt_S02 +loss_cycle_S02
        if self.opt.frequency:
            self.loss_G_S02 += self.loss_F_S2 * self.opt.lambda_fre
        if self.isTrain:
            self.loss_G_S02 += self.criterionGAN(self.netD_S2(self.fake_S02["generation"]), True)*self.opt.lambda_D
            # self.loss_G_S02 += self.criterionCycle(self.fake_S02, self.real_S2)
            # self.loss_G_S02.backward()
        pass
    
    def cal_param_G_S20(self):
        self.loss_G_S20 = torch.tensor(float(0)).to(self.device)
        # if not self.opt.use_inception:
        idt_S20= self.netG_S20(self.real_S0,self.opt.slope)
        # else:
        #     idt_S20,_ = self.netG_S20(self.real_S0,self.opt.slope)
        loss_idt_S20 = self.criterionIdt(idt_S20["generation"], self.real_S0) * self.opt.lambda_A * self.opt.lambda_idt
        loss_cycle_S20 = self.criterionCycle(self.rec_S20, self.real_S0) * self.opt.lambda_A
        
        self.loss_G_S20 += loss_cycle_S20 + loss_idt_S20
        if self.opt.frequency:
            self.loss_G_S20 += self.loss_F_S20 * self.opt.lambda_fre
        if self.isTrain:
            self.loss_G_S20 += self.criterionGAN(self.netD_S0(self.fake_S20["generation"]), True)*self.opt.lambda_D
            # self.loss_G_S20 += self.criterionCycle(self.fake_S20, self.real_S0)
            # self.loss_G_S20.backward()
        pass
    
    def cal_param_G_S0D(self):
        self.loss_G_S0D = torch.tensor(float(0)).to(self.device)
        if self.opt.D_generation == "source":
            # if not self.opt.use_inception:
            self.idt_S0D = self.netG_S0D(self.real_SD,self.opt.slope)
            # else:
            #     idt_S0D,_ = self.netG_S0D(self.real_SD,self.opt.slope)
            loss_idt_S0D = self.criterionIdt(self.idt_S0D["generation"], self.real_SD) * self.opt.lambda_B * self.opt.lambda_idt
            loss_cycle_S0D = self.criterionCycle(self.rec_S0D, self.real_SD) * self.opt.lambda_B
            self.loss_G_S0D += loss_idt_S0D +loss_cycle_S0D
            if self.opt.frequency:
                self.loss_G_S0D += self.loss_F_SD * self.opt.lambda_fre
        elif self.opt.D_generation == "mergeS12":
            if self.opt.frequency:
                self.loss_G_S0D += self.loss_F_S12D * self.opt.lambda_fre
        if self.isTrain:
            if self.opt.D_generation == "source":
                self.loss_G_S0D += self.criterionGAN(self.netD_SD(self.fake_S0D["generation"]), True)*self.opt.lambda_D
                # self.loss_G_S0D += self.criterionCycle(self.fake_S0D, self.real_SD)
            elif self.opt.D_generation == "mergeS12":
                self.loss_G_S0D += self.criterionGAN(self.netD_SD(self.fake_S12D["generation"]), True)*self.opt.lambda_D
                # self.loss_G_S0D += self.criterionCycle(self.fake_S12D, self.real_SD)
            # self.loss_G_S0D.backward()
        pass
    
    def cal_param_G_SD0(self):
        self.loss_G_SD0 = torch.tensor(float(0)).to(self.device)
        if self.opt.D_generation == "source":
            # if not self.opt.use_inception:
            idt_SD0= self.netG_SD0(self.real_S0,self.opt.slope)
            # else:
            #     idt_SD0,_ = self.netG_SD0(self.real_S0,self.opt.slope)
            loss_idt_SD0 = self.criterionIdt(idt_SD0["generation"], self.real_S0) * self.opt.lambda_A * self.opt.lambda_idt
            loss_cycle_SD0 = self.criterionCycle(self.rec_SD0, self.real_S0) * self.opt.lambda_A
            
            self.loss_G_SD0 += loss_cycle_SD0 + loss_idt_SD0
            if self.opt.frequency:
                self.loss_G_SD0 += self.loss_F_SD0 * self.opt.lambda_fre
            if self.isTrain:
                self.loss_G_SD0 += self.criterionGAN(self.netD_S0(self.fake_SD0["generation"]), True)*self.opt.lambda_D
                # self.loss_G_SD0 += self.criterionCycle(self.fake_SD0, self.real_S0)
            # self.loss_G_SD0.backward()
        pass
    
    def calculate_G_parameters(self):
        # if self.isTrain:
        if self.opt.frequency:
            self.backward_Frequency('1')
            self.backward_Frequency('2')
            if self.opt.D_generation == "source":
                self.backward_Frequency('D')
            elif self.opt.D_generation == "mergeS12":
                fake_S12D_F = get_frequency(self.fake_S12D["generation"].contiguous().cpu().detach().numpy()[:,0,:,:])
                self.loss_F_S12D = self.criterionPredD(fake_S12D_F , self.real_SD_F).to(self.device)
            self.loss_F_S01=self.criterionPredD(self.real_S0_F,self.real_S1_F)
            self.loss_F_S02=self.criterionPredD(self.real_S0_F,self.real_S2_F)
            self.loss_F_S0D=self.criterionPredD(self.real_S0_F,self.real_SD_F)

        if self.isTrain:
            self.cal_param_G_S01()
            self.cal_param_G_S10()
            self.cal_param_G_S02()
            self.cal_param_G_S20()
            self.cal_param_G_S0D()
            self.cal_param_G_SD0()
            
            self.loss_G = self.loss_G_S01 + self.loss_G_S10
            self.loss_G += self.loss_G_S02 + self.loss_G_S20
            self.loss_G += self.loss_G_S0D + self.loss_G_SD0
            self.loss_G.backward()
 
            # self.cal_score()
            
        pass
    
    def optimize_parameters(self):
        """Calculate losses, gradients, and update network weights; called in every training iteration"""
        # forward
        self.forward()      # compute fake images and reconstruction images.
        
        # G_A and G_B
        self.set_requires_grad([self.netD_S0, self.netD_S1,self.netD_S2, self.netD_SD,], False)  # Ds require no gradients when optimizing Gs
        self.optimizer_G.zero_grad()
        self.calculate_G_parameters()
        self.optimizer_G.step()

        self.set_requires_grad([self.netD_S0, self.netD_S1,self.netD_S2, self.netD_SD,], True)
        
        self.optimizer_D.zero_grad()
        self.backward_D_S0()
        self.backward_D_S1()
        self.backward_D_S2()
        self.backward_D_SD()
        self.optimizer_D.step()
        
    def cal_score(self):
        if self.opt.frequency:
            self.calculate_G_parameters()

        pass
    def iter_end(self):
        records = {}
        if self.opt.frequency:
            
            records['loss_F_S1'] = self.loss_F_S1
            records['loss_F_S01'] = self.loss_F_S01
            records['loss_F_S2'] = self.loss_F_S2
            records['loss_F_S02'] = self.loss_F_S02
            if self.opt.D_generation == "source":
                records['loss_F_SD'] = self.loss_F_SD
                records['loss_F_S0D'] = self.loss_F_S0D
            elif self.opt.D_generation == "mergeS12":
                records['loss_F_S12D'] = self.loss_F_S12D

        if self.isTrain:
            records["loss_G"]=self.loss_G
            records["loss_G_S01"]=self.loss_G_S01
            records["loss_G_S10"]=self.loss_G_S10
            records["loss_G_S02"]=self.loss_G_S02
            records["loss_G_S20"]=self.loss_G_S20
            records["loss_G_S0D"]=self.loss_G_S0D
            records["loss_G_SD0"]=self.loss_G_SD0
            
            records["loss_D_S0"]=self.loss_D_S0
            records["loss_D_S1"]=self.loss_D_S1
            records["loss_D_S2"]=self.loss_D_S2
            records["loss_D_SD"]=self.loss_D_SD
                
        else:

            records["fake_S01"] = self.fake_S01["generation"]
            records["fake_S02"] = self.fake_S02["generation"]
            # records["idt_S01"] = self.idt_S01
            # records["idt_S02"] = self.idt_S02
            # records["fake_S10"] = self.fake_S10
            # records["fake_S20"] = self.fake_S20
            # records["rec_S10"] = self.rec_S10
            # records["rec_S20"] = self.rec_S20
            records["real_S0"] = self.real_S0
            records["real_S1"] = self.real_S1
            records["real_S2"] = self.real_S2
            records["real_SD"] = self.real_SD
            if self.opt.use_inception:
                records["loss_FID_S1"] = self.loss_FID_S1
                records["loss_FID_S2"] = self.loss_FID_S2
                records["loss_FID_SD"] = self.loss_FID_SD
                
                records["loss_FID_S01"] = self.loss_FID_S01
                records["loss_FID_S02"] = self.loss_FID_S02
                records["loss_FID_S0D"] = self.loss_FID_S0D
                
            #     records["real_S0_Fea"] = self.real_S0_Fea
            #     records["real_S1_Fea"] = self.real_S1_Fea
            #     records["real_S2_Fea"] = self.real_S2_Fea
            #     records["real_SD_Fea"] = self.real_SD_Fea
            #     records["fake_S1_Fea"] = self.fake_S1_Fea
            #     records["fake_S2_Fea"] = self.fake_S2_Fea 
            #     records["fake_SD_Fea"] = self.fake_SD_Fea
                
            records["image_paths"] = self.image_paths
            if self.opt.D_generation == "source":
                records["fake_S0D"] = self.fake_S0D
                # records["fake_SD0"] = self.fake_SD0
                # records["idt_S0D"] = self.rec_S0D
                # records["rec_SD0"] = self.rec_SD0
            elif self.opt.D_generation == "mergeS12":
                records["fake_S12D"] = self.fake_S12D["generation"]
                # if self.opt.use_inception:
                # records["loss_FID_S0D"] = self.loss_FID_S0D
                # records["loss_FID_S01"] = self.loss_FID_S01
                # records["loss_FID_S02"] = self.loss_FID_S02
                
                # records["loss_FID_SD"] = self.loss_FID_SD
                # records["loss_FID_S1"] = self.loss_FID_S1
                # records["loss_FID_S2"] = self.loss_FID_S2
                
        records = self.TestApi.record(records)
        
        if self.isTrain:
            self.loss_mean_G=records["count_mean_G"]
            self.loss_mean_G_S01=records["count_mean_G_S01"]
            self.loss_mean_G_S10=records["count_mean_G_S10"]
            self.loss_mean_G_S02=records["count_mean_G_S02"]
            self.loss_mean_G_S20=records["count_mean_G_S20"]
            self.loss_mean_G_S0D=records["count_mean_G_S0D"]
            self.loss_mean_G_SD0=records["count_mean_G_SD0"]
            
            self.loss_mean_D_S0=records["count_mean_D_S0"]
            self.loss_mean_D_S1=records["count_mean_D_S1"]
            self.loss_mean_D_S2=records["count_mean_D_S2"]
            self.loss_mean_D_SD=records["count_mean_D_SD"]
        else:
            if self.opt.use_inception:
                self.loss_mean_FID_S1 = records["count_mean_FID_S1"]
                self.loss_mean_FID_S2 = records["count_mean_FID_S2"]
                self.loss_mean_FID_SD = records["count_mean_FID_SD"]
                self.loss_mean_FID_S01 = records["count_mean_FID_S01"]
                self.loss_mean_FID_S02 = records["count_mean_FID_S02"]
                self.loss_mean_FID_S0D = records["count_mean_FID_S0D"]
                self.loss_mean_DFID_S1 = records["count_mean_S1disFID"]
                self.loss_mean_DFID_S2 = records["count_mean_S2disFID"]
                self.loss_mean_DFID_SD = records["count_mean_SDdisFID"]
            # self.loss_mean_S012D_SSIM = records["count_mean_S012D_SSIM.avg
            
            if self.opt.D_generation == "source":
                self.loss_mean_PSNR_SD = records["count_mean_PSNR_SD"]
                self.loss_mean_SSIM_SD = records["count_mean_SSIM_SD"]
            elif self.opt.D_generation == "mergeS12":
               self.loss_mean_PSNR_S12D = records["count_mean_PSNR_S12D"]
               self.loss_mean_SSIM_S12D = records["count_mean_SSIM_S12D"]
                    
            self.loss_mean_PSNR_S1 = records["count_mean_PSNR_S1"]
            self.loss_mean_PSNR_S2 = records["count_mean_PSNR_S2"]
            self.loss_mean_S01_PSNR = records["count_mean_S01_PSNR"]
            self.loss_mean_S02_PSNR = records["count_mean_S02_PSNR"]
            self.loss_mean_S0D_PSNR = records["count_mean_S0D_PSNR"]

            self.loss_mean_SSIM_S1 = records["count_mean_SSIM_S1"]
            self.loss_mean_SSIM_S2 = records["count_mean_SSIM_S2"]
            self.loss_mean_S01_SSIM = records["count_mean_S01_SSIM"]
            self.loss_mean_S02_SSIM = records["count_mean_S02_SSIM"]
            self.loss_mean_S0D_SSIM = records["count_mean_S0D_SSIM"]
            
        if self.opt.frequency:
            self.loss_mean_F_S1 = records["count_mean_F_S1"]
            self.loss_mean_F_S01 = records["count_mean_F_S01"]
            self.loss_mean_F_S2 = records["count_mean_F_S2"]
            self.loss_mean_F_S02 = records["count_mean_F_S02"]
            if self.opt.D_generation == "source":
                self.loss_mean_F_SD = records["count_mean_F_SD"]
                self.loss_mean_F_S0D = records["count_mean_F_S0D"]
            elif self.opt.D_generation == "mergeS12":
                self.loss_mean_F_S12D = records["count_mean_F_S12D"]
        self.fake_S01 = self.fake_S01["generation"]
        self.rec_S10 = self.rec_S10["generation"]
        self.fake_S02 = self.fake_S02["generation"]
        self.rec_S20 = self.rec_S20["generation"]
        self.fake_S0D = self.fake_S0D["generation"]
        self.rec_SD0 = self.rec_SD0["generation"]        
        
        if self.isTrain:
            self.fake_S10 = self.fake_S10["generation"]
            self.rec_S01 = self.rec_S01["generation"]
            self.fake_S20 = self.fake_S20["generation"]
            self.rec_S02 = self.rec_S02["generation"]
            self.fake_SD0 = self.fake_SD0["generation"]
            self.rec_S0D = self.rec_S0D["generation"]
        pass
    
    def epoch_end(self):
        records = self.TestApi.statistic()
        if not self.isTrain:
            self.TOP10p_mean_SSIM_S1 = records["TOP10p_mean_SSIM_S1"]
            self.TOP50p_mean_SSIM_S1 = records["TOP50p_mean_SSIM_S1"]
            self.TOP100p_mean_SSIM_S1 = records["TOP100p_mean_SSIM_S1"]
            
            self.TOP10p_mean_SSIM_S2 = records["TOP10p_mean_SSIM_S2"]
            self.TOP50p_mean_SSIM_S2 = records["TOP50p_mean_SSIM_S2"]
            self.TOP100p_mean_SSIM_S2 = records["TOP100p_mean_SSIM_S2"]
            
            self.TOP10p_mean_SSIM_SD = records["TOP10p_mean_SSIM_SD"]
            self.TOP50p_mean_SSIM_SD = records["TOP50p_mean_SSIM_SD"]
            self.TOP100p_mean_SSIM_SD = records["TOP100p_mean_SSIM_SD"]
        
                  
    def val_save(self,key=None):
        # self.val_best = (self.count_mean_SSIM_S1.avg + self.count_mean_SSIM_S2.avg + self.count_mean_SSIM_SD.avg)
        ans = False
        if key is not None and key != '':
            _list = ['S1', 'S2', 'SD']
            self.model_names = []
            
            for modality in _list:
                tmp = key.split("_")
                tmp = '_'.join(tmp+[modality])
                val_best = getattr(self,tmp)
                if tmp in self.last_val_best:
                    if self.last_val_best[tmp] < val_best:
                        self.last_val_best[tmp] = val_best
                        if 'SD' in tmp:
                            self.model_names += ['G_S0D', 'G_SD0']
                        elif 'S1' in tmp:
                            self.model_names += ['G_S01', 'G_S10',]
                        elif 'S2' in tmp:
                            self.model_names += ['G_S02', 'G_S20',]
                        ans = True
                else:
                    self.last_val_best[tmp] = val_best
                    if 'SD' in tmp:
                        self.model_names += ['G_S0D', 'G_SD0']
                    elif 'S1' in tmp:
                        self.model_names += ['G_S01', 'G_S10',]
                    elif 'S2' in tmp:
                        self.model_names += ['G_S02', 'G_S20',]
                    ans = True
        return ans
    

    # def set_validation_parameter(self):
    #     if self.opt.val_param is not '':
    #         # self.val_param = ['Top10meanSSIM','Top50meanSSIM','Top100meanSSIM','Mse']
    #         self.val_param = self.opt.val_param.split(',')
    #     else:
    #         self.val_param = [None]
    
    # def get_validation_parameter(self):
    #     # self.val_param = ['Top10meanSSIM','Top50meanSSIM','Top100meanSSIM','Mse']
    #     ans = []
    #     for key in self.val_param:
    #         if key in self.__dir__():
    #             ans.append(key)
    #     if len(ans) == 0:
    #         return None
    #     else:
    #         return ans
        
        
        # if self.opt.isTrain:
        #     self.epoch+=1
            # self.fac = self.epoch/(self.nepoch+1)
